package org.example.myleet.p745;

import java.util.HashMap;
import java.util.Map;

public class WordFilter {

    private final MyTrieTree trieTree;

    public WordFilter(String[] words) {
        trieTree = new MyTrieTree();
        for (int i = 0; i < words.length; ++i) {
            trieTree.add(words[i], i);
        }
    }

    public int f(String pref, String suff) {
        Map<String, Integer> wordMapId = trieTree.findWordIdByPrefixAndSuffix(pref, suff);
        int id = -1;
        for (Map.Entry<String, Integer> entry : wordMapId.entrySet()) {
            id = Math.max(id, entry.getValue());
        }
        return id;
    }

    static class MyTrieTree {
        TrieNode root;

        public MyTrieTree() {
            root = new TrieNode('0');
        }

        public void add(String word, int id) {
            TrieNode cur = root;
            for (int i = 0; i < word.length(); ++i) {
                char c = word.charAt(i);
                int idx = c - 'a';
                if (null == cur.children[idx]) {
                    cur.children[idx] = new TrieNode(c);
                }
                cur = cur.children[idx];
            }
            cur.end = true;
            cur.word = word;
            if (null == cur.id) {
                cur.id = id;
            } else {
                cur.id = Math.max(cur.id, id);
            }
        }

        public Map<String, Integer> findWordIdByPrefixAndSuffix(String prefix, String suffix) {
            Map<String, Integer> wordMapId = new HashMap<>();
            TrieNode cur = root;
            for (int i = 0; i < prefix.length(); ++i) {
                char c = prefix.charAt(i);
                int idx = c - 'a';
                if (null == cur.children[idx]) {
                    return new HashMap<>();
                }
                cur = cur.children[idx];
            }
            dfs(cur, wordMapId, suffix);
            return wordMapId;
        }

        private void dfs(TrieNode node, Map<String, Integer> wordMapId, String suffix) {
            if (node.end && node.word.endsWith(suffix)) {
                wordMapId.put(node.word, node.id);
            }
            for (TrieNode child : node.children) {
                if (null != child) {
                    dfs(child, wordMapId, suffix);
                }
            }
        }
    }

    static class TrieNode {
        TrieNode[] children;
        boolean end;
        char c;
        String word;
        Integer id;

        public TrieNode(char c) {
            this.c = c;
            end = false;
            children = new TrieNode[26];
            word = null;
            id = null;
        }
    }
}
